# -*- coding: utf-8 -*-

# Copyright 2009 David Lee Miller
#
# This file is part of POSSM (Python Open Source Scanning Microscopy).
#
# POSSM is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# POSSM is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License         
# along with POSSM.  If not, see <http://www.gnu.org/licenses/>.
#                                                                           
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

from display import Scanning2,NotebookSpec,NotebookPage,FakePkg,MyNavigationToolbar
import wx.lib.flatnotebook as wxf
from Scientific.IO import NetCDF
from client import copy,deepcopy,wx,time,ws
from client import n #the numpy package
from client import matplotlib,cm,Toolbar,FigureCanvasWxAgg,NavigationToolbar2WxAgg,Figure
import string



class Datapkg(object):
	data = None

class FakeQueue(object):
	'''
	A class with some attributes of a queue so that
	it can be used in place.
	'''
	def put(self,data):
		pass
	def get(self):
		pass

class Scanview(wx.Frame,Scanning2):
	'''
	Contains utilities for viewing saved data
	in NetCDF format. 
	'''
	dataq = FakeQueue()
	specq = FakeQueue()
	disp_index = 0
	ASPECT_W = ws.ASPECT_W
	ASPECT_H = ws.ASPECT_H
	ASPECT_SIZE = ws.ASPECT_SIZE
	TITLE_SIZE = ws.TITLE_SIZE
	AXES_SIZE = ws.AXES_SIZE
	DEFAULT_COLOR = ws.DEFAULT_COLOR
	POINT_TYPE = ws.POINT_TYPE
	EXTENSION = ws.EXTENSION
	images = ['image','imagewiv','imagewiz','imagewzv','multibias','continuous']
	spectroscopy = ['imagewiv','imagewiz','imagewzv']
	single = ['singleiv','singlezv','singleiz']
	def __init__(self, *args, **kwds):
		kwds["style"] = wx.DEFAULT_FRAME_STYLE
		homefrm=kwds['homefrm']
		self.editor=homefrm.editor
		wx.Frame.__init__(self, args[0], -1, "", pos=(0, 0), size=ws.DEFAULT_SIZE)
		bookStyle = wxf.FNB_NODRAG
		self.nb = wxf.FlatNotebook(self, wx.ID_ANY, style=bookStyle)
		self.nb.Bind(wxf.EVT_FLATNOTEBOOK_PAGE_CHANGED,self.OnChange)
		self.panel = wx.Panel(self,-1)
		self.fig = Figure((self.ASPECT_W,self.ASPECT_H), self.ASPECT_SIZE)

		self.canvas = FigureCanvasWxAgg(self.panel, -1, self.fig)
		self.toolbar = MyNavigationToolbar(self,self.canvas, True, specselect=False)
		self.toolbar.Realize()
		#self.canvas.SetForegroundColour(wx.WHITE)
		self.add_page = wx.Button(self.panel,-1,'Add Page')
		self.add_page.Bind(wx.EVT_BUTTON, self.OnAdd)
		if not hasattr(self, 'subplot'):
                	self.subplot = self.fig.add_subplot(111)
		self.nb.AddPage(self.panel,"Preview")
		Scanning2.__init__(self,self.nb,dataq=self.dataq,specq=self.specq,homefrm=homefrm)
		self.t.Stop()
		self._directory()
		self.__set_properties()
		self.__do_layout()
		self.SetTitle("Viewer")
	
	def __do_layout(self):
		'''
		Organizes the wx widgets on the frame.
		'''
		sizer_1 = wx.FlexGridSizer(1, 1, 0, 0)
		sizer_2 = wx.FlexGridSizer(1, 2, 0, 0)
		sizer_3 = wx.FlexGridSizer(2, 1, 0, 0)
		sizer_4 = wx.FlexGridSizer(2, 1, 0, 0)
		fbbox = wx.StaticBox(self.panel, -1, 'File Browser')
		sizer2b = wx.StaticBoxSizer(fbbox,wx.VERTICAL)
		sizer_4.Add(self.file_browser, 0, wx.ALL, 5)
		sizer_4.Add(self.add_page, 0, wx.ALL, 5)
		sizer2b.Add(sizer_4, 0, wx.ALL, 5)
		#sizer2b.Add(self.new_page, 0, wx.ALL, 5)
		#grid_sizer_1.AddGrowableRow(0)
		#grid_sizer_1.AddGrowableCol(1)
		#sizer_1.Add(sizer2b, 1, wx.EXPAND|wx.ALL, 5)
		sizer_1.Add(self.nb, 1, wx.EXPAND|wx.ALL, 5)
		sizer_2.Add(sizer2b, 1, wx.EXPAND|wx.ALL, 5)
		sizer_3.Add(self.canvas, 1, wx.EXPAND|wx.ADJUST_MINSIZE|wx.ALL, 5)
		sizer_3.Add(self.toolbar, 1, wx.GROW, 0)
		sizer_2.Add(sizer_3, 1, wx.EXPAND|wx.ALL, 0)
		self.SetAutoLayout(True)
		self.SetSizer(sizer_1)
		self.SetSizer(sizer_2)
		self.Layout()

	def __set_properties(self):
		'''
		Sets the properties of the wx widgets used.
		'''
		width,height = ws.DEFAULT_SIZE
		height -= 150
		width -= 50
		fbwidth = 200
		self.file_browser.SetMinSize((fbwidth, height-80))
		self.nb.SetMinSize((width,height))
		self.canvas.SetMinSize((width-fbwidth,height-50))
		self.add_page.SetMinSize((fbwidth,35))
	
	def _directory(self):
		'''
		Sets the working directory for opening data.
		'''
		self.path = ws.BASE_DIRECTORY+ws.DATA_DIRECTORY+'/'+time.strftime('%Y')
		self.file_browser = wx.GenericDirCtrl(self.panel, -1, size=(200,250), style=0)
		file_ctrl = self.file_browser.GetTreeCtrl()
		file_ctrl.Bind(wx.EVT_TREE_SEL_CHANGED, self.OnSelect)
		self.file_browser.SetPath(self.path)
	
	def do_show(self):
		self.Show()
	
	def OnChange(self,event):
		'''
		Toggles events when different tabs are selected. This
		toggles the X on the tab so that the preview tab cannot
		be closed.
		'''
		if self.nb.GetPageText(self.nb.GetSelection()) == 'Preview':
			style = self.nb.GetWindowStyleFlag()
			style |= wxf.FNB_X_ON_TAB
			style ^= wxf.FNB_X_ON_TAB
			self.nb.SetWindowStyleFlag(style)
		else:
			style = self.nb.GetWindowStyleFlag()
			style |= wxf.FNB_X_ON_TAB
			self.nb.SetWindowStyleFlag(style)
	
	def OnSelect(self, event):
		'''
		Checks the extension of the file before adding it
		to the preview window.
		'''
		leng = len(self.EXTENSION)
		if self.file_browser.GetPath()[-leng:] == self.EXTENSION:
			self.OnPreview(None)
		elif self.file_browser.GetPath()[-leng:] == string.upper(self.EXTENSION):
			self.OnPreview(None)
	
	def OnPreview(self,event):
		'''
		Reads data from the file and quickly displays it. 
		'''
		path = self.file_browser.GetPath()
		
		if path[-3:] == '.nc' or path[-3:] == '.NC':
			netfile = NetCDF.NetCDFFile(path,'r')
			self.datapkg = FakePkg()
			if hasattr(netfile,'TYPE'):
				self.TYPE = netfile.TYPE
			elif 'image_data' in netfile.variables.keys():
				self.TYPE = 'NIST'
			else: self.TYPE = None
			self.subplot.cla()
			if self.TYPE in self.images:
				self.datapkg.IMAGEDATA = netfile.variables['IMAGEDATA'].getValue()
				#if not hasattr(self,'im'):
				self.im=self.subplot.imshow(n.array(self.datapkg.IMAGEDATA[0][0]),
					cmap=cm.gray, origin='lower')
				self.subplot.set_ylabel('Y (pixels)')
				self.subplot.set_xlabel('X (pixels)')
				self.subplot.set_title('Preview')
			elif self.TYPE in self.single:
				#if hasattr(self,'im'): self.im.remove()
				subtype=self.TYPE[-2:]
				data=[]
				spec=netfile.variables[subtype].getValue()
				xspec=netfile.variables['x'+subtype].getValue()
				yspec=netfile.variables['y'+subtype].getValue()
				if not subtype=='iz': xvals=netfile.variables['V'].getValue()
				else: xvals=netfile.variables['Z_VALS'].getValue()
				data=[xvals,spec[0,0,0,:],xvals,spec[0,1,0,:]]
				self.subplot.plot(data[0],data[1],'-b',data[2],data[3],'-r')
				self.subplot.set_ylabel(string.upper(subtype[0]))
				self.subplot.set_xlabel(string.upper(subtype[1]))
				self.subplot.set_title('Preview')
				#print dir(self.subplot)
			elif self.TYPE == 'NIST':
				self.datapkg.IMAGEDATA = [[netfile.variables['image_data'].getValue()]]
				#if not hasattr(self,'im'):
				self.im=self.subplot.imshow(n.array(self.datapkg.IMAGEDATA[0][0]),
					cmap=cm.gray, origin='lower')
				self.subplot.set_ylabel('Y (pixels)')
				self.subplot.set_xlabel('X (pixels)')
				self.subplot.set_title('Preview')
			self.canvas.draw()
	
	def OnAdd(self,event):
		self.OnSet()
	
	def OnSet(self):
		'''
		Takes the image in preview and gives it its own tab in the notebook.
		It starts by filling the data package from the netcdf file. Then adds
		it to the notebook, creates the image or plot and flags the data as
		'complete'.
		'''
		path = self.file_browser.GetPath()
		
		# Opens the file and creates the container for the data
		if path[-3:] == '.nc' or path[-3:] == '.NC':
			self.dic = {}
			netfile = NetCDF.NetCDFFile(path,'r')
			self.netfile = netfile
			self.datapkg = FakePkg()
			if hasattr(netfile,'TYPE'):
				self.TYPE = netfile.TYPE
				self.datapkg.NAME = netfile.NAME
			elif 'image_data' in netfile.variables.keys():
				self.TYPE = 'NIST'
			else: self.TYPE = None
			self.disp_index += 1
			self.datapkg.PATH = path
			self.datapkg.SCANSTAMP = copy(self.disp_index)
			
			# Add the data properties to the data package
			if self.TYPE in self.images:
				for i in netfile.variables.keys():
					if hasattr(netfile.variables[i],'unit'): self.datapkg.__setattr__(i,(n.array(netfile.variables[i].getValue())[0],netfile.variables[i].unit))
					else: self.datapkg.__setattr__(i,netfile.variables[i].getValue())
				self.datapkg.IMAGEDATA=n.array(self.datapkg.IMAGEDATA)
				self.datapkg.SCAN = netfile.SCAN_ID
				self.datapkg.TYPE = netfile.TYPE
				self.datapkg.SCANSTAMP = netfile.SCAN_ID
				self.datapkg.SHAPE = self.datapkg.IMAGEDATA[0][0].shape
				self.datapkg.EXTENT = (None,None,None,None)
				if 'BIAS_STEPS' in netfile.variables:
					self.datapkg.BIAS_STEPS=[]
					BIAS_STEPS = netfile.variables['BIAS_STEPS'].getValue()
					for i in range(len(BIAS_STEPS)): self.datapkg.BIAS_STEPS.append((BIAS_STEPS[i],))
				fig = Figure((ws.IM_ASPECT_W,ws.IM_ASPECT_H), ws.IM_ASPECT_SIZE)
				number = len(self.pages)+1
				self.pages[str(self.datapkg.SCANSTAMP)] = NotebookPage(self,self.nb,number,editor=self.editor)
				page=self.pages[str(self.datapkg.SCANSTAMP)]
				page.addToNotebook(str(self.datapkg.SCANSTAMP))
				page.addPlot(fig,self.datapkg)
				page.lineN = self.datapkg.SHAPE[0]
				if self.TYPE in self.spectroscopy:
					subtype = self.TYPE[-2:]
					data={}
					spec=netfile.variables[subtype].getValue()
					xspec=netfile.variables['x'+subtype].getValue()
					yspec=netfile.variables['y'+subtype].getValue()
					if not subtype=='iz': xvals=netfile.variables['V'].getValue()
					else: xvals=netfile.variables['Z_VALS'].getValue()
					for i,j in enumerate(xspec):
						loc=str((int(yspec[i]),int(xspec[i])))
						data[loc]=[[[n.array(xvals),n.array(spec[0,0,i,:])],[n.array(xvals),n.array(spec[0,1,i,:])]]]
					self.datapkg.SPEC_MASK=n.array(self.datapkg.SPEC_MASK)
					exec 'self.datapkg.__setattr__(\''+string.upper(subtype)+'\',data)'
					#page.dispSpec(n.array(self.datapkg.SPEC_MASK))
				page._set_complete(self.datapkg)
				netfile.close()
			elif self.TYPE in self.single:
				self.datapkg.SCAN = netfile.SCAN_ID
				self.datapkg.TYPE = netfile.TYPE
				self.datapkg.SCANSTAMP = netfile.SCAN_ID
				self.datapkg.SHAPE = (10,10)
				self.datapkg.EXTENT = (None,None,None,None)
				self.datapkg.GAIN_ZHV = [netfile.variables['GAIN_ZHV'].getValue(),'']
				self.datapkg.GAIN_ZOUT = [netfile.variables['GAIN_ZOUT'].getValue(),'']
				self.datapkg.GAIN_SERVO = [netfile.variables['GAIN_SERVO'].getValue(),'']
				subtype = self.TYPE[-2:]
				#print netfile.variables
				data={}
				spec=netfile.variables[subtype].getValue()
				xspec=netfile.variables['x'+subtype].getValue()
				yspec=netfile.variables['y'+subtype].getValue()
				if not subtype=='iz': xvals=netfile.variables['V'].getValue()
				else: xvals=netfile.variables['Z_VALS'].getValue()
				for i,j in enumerate(xspec):
					loc=str((int(yspec[i]),int(xspec[i])))
					data[loc]=[[[n.array(xvals),n.array(spec[0,0,i,:])],[n.array(xvals),n.array(spec[0,1,i,:])]]]
				exec 'self.datapkg.'+string.upper(subtype)+'=data'
				fig = Figure((ws.IM_ASPECT_W,ws.IM_ASPECT_H), ws.IM_ASPECT_SIZE)
				number = len(self.pages)
				self.pages[str(self.datapkg.SCANSTAMP)] = NotebookSpec(self,self.nb,number)
				self.pages[str(self.datapkg.SCANSTAMP)].addToNotebook(str(self.datapkg.SCANSTAMP))
				self.pages[str(self.datapkg.SCANSTAMP)].addPlot(fig,self.datapkg,subtype)
				netfile.close()
			elif self.TYPE == 'NIST':
				for i in netfile.variables.keys():
					if hasattr(netfile.variables[i],'unit'): self.datapkg.__setattr__(i,(n.array(netfile.variables[i].getValue())[0],netfile.variables[i].unit))
					else: self.datapkg.__setattr__(i,netfile.variables[i].getValue())
				self.datapkg.IMAGEDATA=[[n.array(self.datapkg.image_data)]]
				self.datapkg.SCAN = 'NISTexperimental'
				self.datapkg.TYPE = 'NIST'
				#self.datapkg.SCANSTAMP = '000001'
				self.datapkg.SHAPE = self.datapkg.IMAGEDATA[0][0].shape
				self.datapkg.EXTENT = (None,None,None,None)
				fig = Figure((ws.IM_ASPECT_W,ws.IM_ASPECT_H), ws.IM_ASPECT_SIZE)
				number = len(self.pages)+1
				self.pages[str(self.datapkg.SCANSTAMP)] = NotebookPage(self,self.nb,number,editor=self.editor)
				page=self.pages[str(self.datapkg.SCANSTAMP)]
				page.addToNotebook(str(self.datapkg.SCANSTAMP))
				page.addPlot(fig,self.datapkg)
				page.lineN = self.datapkg.SHAPE[0]
				page._set_complete(self.datapkg)
				netfile.close()